package small.common.limit.leakyBucket;

import small.common.core.util.TimeUtil;
import small.common.limit.TrafficShaping;
import small.common.limit.TrafficShapingNode;

import java.util.concurrent.atomic.AtomicLong;

/**
 * 漏桶算法
 *
 * @author jinzhou
 * @data 2023/4/7 17:51
 */
public class LeakyBucketTrafficShaping implements TrafficShaping {

    /**
     * 最大等待时间
     */
    private final int maxQueueingTimeMs;

    /**
     * 阀值
     */
    private final double count;

    /**
     * 上次通过该资源的时间戳
     */
    private final AtomicLong latestPassedTime = new AtomicLong(-1);

    public LeakyBucketTrafficShaping(int timeOut, double count) {
        this.maxQueueingTimeMs = timeOut;
        this.count = count;
    }

    @Override
    public boolean canPass(TrafficShapingNode node, long acquireCount) {
        if (acquireCount <= 0) {
            return true;
        }
        if (count <= 0) {
            return false;
        }
        //当前时间
        long currentTimeMillis = TimeUtil.currentTimeMillis();
        //两次请求之间时间间隔(ms)
        long costTime = Math.round(1.0 * (acquireCount) / count * 1000);
        //说明上次请求已经很久
        if ((latestPassedTime.get() + costTime) <= currentTimeMillis) {
            latestPassedTime.set(currentTimeMillis);
            return true;
        } else {
            //等待时间
            long waitTime = latestPassedTime.get() + costTime - TimeUtil.currentTimeMillis();
            //等待时间超过最大时间
            if (waitTime > maxQueueingTimeMs) {
                return false;
            } else {
                long oldTime = latestPassedTime.addAndGet(costTime);
                try {
                    waitTime = oldTime - TimeUtil.currentTimeMillis();
                    if (waitTime > maxQueueingTimeMs) {
                        latestPassedTime.addAndGet(-costTime);
                        return false;
                    }
                    //睡眠一段时间
                    if (waitTime > 0) {
                        Thread.sleep(waitTime);
                    }
                    return true;
                } catch (InterruptedException e) {
                }
            }

        }

        return false;
    }
}
